# coding=utf-8

import os
import cv2
import numpy as np
from PIL import Image
from pylab import *

class IMBase():
    def __init__(self):
        self.path = "./"

class IMSift(IMBase):
    def process_image(self, imagename, resultname, params="--edge-thresh 10 --peak-thresh 5"):
        if imagename[-3:] != "pgm":
            im = Image.open(imagename).convert('L')
            im.save(self.path + 'tmp.pgm')
            imagename = "tmp.pgm"
        
        cmmd = str("sift " + imagename + " --output=" + resultname + 
            " " + params)
        print(cmmd)
        os.system(cmmd)
        print('processed ' + imagename, ' to ', resultname)

    def process_image_demo(self, imagename, resultname):
        self.process_image(imagename, resultname)

    def read_features_from_file(self, filename):
        f = np.loadtxt(filename)
        return f[:, :4], f[:, :4]

    def write_features_to_file(self, filename, locs, desc):
        np.savetxt(filename, np.hstack((locs, desc)))
    
    def plot_futures(self, im, locs, circle=False):
        def draw_circle(c, r):
            t = np.arange(0, 1.01, .01)*2*np.pi
            x = r *np.cos(t) + c[0]
            y = r * np.sin(t) + c[1]
            plot(x, y, 'b', linewidth=2)

        if circle:
            for p in locs:
                draw_circle(p[:2], [2])
        else:
            plot(locs[:, 0], locs[:, 1], 'ob')
        
        axis('off')

    def demo_sift(self, filename):
        iml = array(Image.open(filename).convert('L'))
        self.process_image(filename, 'output.sift')
        l1, d1 = self.read_features_from_file('output.sift')

        figure()
        gray()
        self.plot_futures(iml, l1, circle=True)
        show()
    
    def match(self, desc1, desc2):
        desc1 = array([d/linalg.norm(d) for d in desc1])
        desc2 = array(d/linalg.norm(d) for d in desc2)

        dist_ratio = 0.6
        desc1_size = desc1.shape

        matchscores = zeros((desc1_size[0], 1), 'int')
        desc2t = desc2.T
        for i in range(desc1_size[0]):
            dotprods = dot(desc1[i, :], desc2t)
            dotprods = 0.9999 * dotprods
            indx = argsort(arccos(dotprods))

            if arccos(dotprods)[indx[0]] < dist_ratio * arccos(dotprods)[indx[1]]:
                matchscores[i] = int(indx[0])

        return matchscores
    
    def match_twosides(self, desc1, desc2):
        matches_12 = self.match(desc1, desc2)
        matches_21 = self.match(desc2, desc1)

        ndx_12 = matches_12.nonzero()[0]

        for n in ndx_12:
            if matches_21[int(matches_12[n])] != n:
                matches_12[n] = 0
            
        return matches_12

    def cv_sift(self, file):
        img = cv2.imread(file)
        gray= cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)

        sift = cv2.xfeatures2d.SIFT_create()
        kp = sift.detect(gray,None)

        img=cv2.drawKeypoints(gray,kp, img)
        cv2.imshow('x', img)
        cv2.waitKey(0)
